import { useCallback, useEffect, useId, useRef, useState } from "react";
import useSWR from "swr";
import { nanoid } from "./nanoid";

// 提交对话，处理返回的流式数据
const streamResponse = async (
  chatRequest,
  action,
  mutate,
  messagesRef,
  onFinish,
  genId,
) => {
  // 更新状态，显示最新的 message。怎么实现的？
  const previousMessages = messagesRef.current;

  mutate(chatRequest.messages, false);

  const replyId = genId();
  const createdAt = new Date();

  let responseMessage = {
    id: replyId,
    createdAt,
    content: "",
    role: "assistant",
  };

  // promise 是一个 promise 对象，包含了 content, ui, next 三个属性
  async function readRow(promise) {
    const { content, ui, next } = await promise;

    console.log("content", content);

    responseMessage["content"] = content.value;
    responseMessage["ui"] = await ui;

    // 在 messages 的最后一条"role"是 "assistant" 的 content 增加 content 内容
    // 如果最后一条是 "user" 的 content，就增加一条新的 "assistant" 的 content
    const lastMessage = messagesRef.current[messagesRef.current.length - 1];

    if (lastMessage.role === "assistant") {
      lastMessage.content += content.value;
      mutate([...messagesRef.current], false);
    } else {
      mutate([...messagesRef.current, { ...responseMessage }], false);
    }

    if (next) {
      await readRow(next);
    }
  }

  try {
    const promise = action({
      messages: chatRequest.messages,
    });

    console.log("promise", promise);

    await readRow(promise);
  } catch (e) {
    // 如果出错，就恢复之前的消息
    mutate(previousMessages, false);
    throw e;
  }

  if (onFinish) {
    onFinish(responseMessage);
  }

  return responseMessage;
};

// 对话的 hook 函数, action 是server action函数, onFinish 是对话结束之后的处理,onError 是处理错误的函数
export function useChat({ action, onFinish, onError, genId = nanoid }) {
  const id = useId();

  // 保存一个空的数组作为最开始的初始化消息
  // FB 是 fallback 的意思, msg 是 message 的意思
  const [initMsgFB] = useState([]);

  // 保存 chat 状态在 SWR。SWR 可以缓存请求，有更新的时候才会发起调用
  // data: messages 是一个对象解构赋值语句。它将 useSWR 函数返回的对象的 data 属性解构为一个名为 messages 的变量。
  // https://swr.vercel.app/docs/api
  // 使用 [id,"message"] 作为 key，这样就可以根据 id 来区分不同的内容
  const { data: messages, mutate } = useSWR([id, "messages"], null, {
    fallbackData: initMsgFB,
  });

  // 同样的，保存 loading 状态在 SWR
  const { data: isLoading = false, mutate: mutateLoading } = useSWR(
    [id, "loading"],
    null,
  );

  // 错误处理
  const { data: error = undefined, mutate: setError } = useSWR(
    [id, "error"],
    null,
  );

  // 保持对最新消息的引用
  const messagesRef = useRef(messages || []);

  useEffect(() => {
    messagesRef.current = messages || [];
  }, [messages]);

  // 取消当前调用的引用
  const abortRef = useRef(null);

  const triggerRequest = useCallback(
    async (chatRequest) => {
      try {
        mutateLoading(true);
        setError(undefined);

        const abortController = new AbortController();
        abortRef.current = abortController;

        await streamResponse(
          chatRequest,
          action,
          mutate,
          messagesRef,
          onFinish,
          genId,
        );

        abortRef.current = null;
      } catch (err) {
        // Ignore abort errors as they are expected.
        if (err.name === "AbortError") {
          abortRef.current = null;
          return null;
        }

        if (onError && err instanceof Error) {
          onError(err);
        }

        setError(err);
      } finally {
        mutateLoading(false);
      }
    },
    [
      id,
      mutate,
      abortRef,
      messagesRef,
      action,
      onError,
      onFinish,
      setError,
      abortRef,
    ],
  );

  const append = useCallback(
    async (message) => {
      if (!message.id) {
        message.id = genId();
      }

      const chatRequest = {
        messages: messagesRef.current.concat(message),
      };

      return triggerRequest(chatRequest);
    },
    [triggerRequest, genId],
  );

  const reload = useCallback(
    async ({ options } = {}) => {
      if (messagesRef.current.length === 0) return null;

      // Remove last assistant message and retry last user message.
      const lastMessage = messagesRef.current[messagesRef.current.length - 1];
      if (lastMessage.role === "assistant") {
        const chatRequest = {
          messages: messagesRef.current.slice(0, -1),
          options,
        };

        return triggerRequest(chatRequest);
      }

      const chatRequest = {
        messages: messagesRef.current,
        options,
      };

      return triggerRequest(chatRequest);
    },
    [triggerRequest],
  );

  const stop = useCallback(() => {
    if (abortRef.current) {
      abortRef.current.abort();
      abortRef.current = null;
    }
  }, []);

  const setMessages = useCallback(
    (messages) => {
      mutate(messages, false);
      messagesRef.current = messages;
    },
    [mutate],
  );

  // Input state and handlers.
  const [input, setInput] = useState("");

  const handleSubmit = useCallback(
    (e) => {
      e.preventDefault();
      if (!input) return;

      append({
        content: input,
        role: "user",
        createdAt: new Date(),
      });
      setInput("");
    },
    [input, append],
  );

  const handleInputChange = (e) => {
    setInput(e.target.value);
  };

  return {
    messages: messages || [],
    error,
    append,
    reload,
    stop,
    setMessages,
    input,
    setInput,
    handleInputChange,
    handleSubmit,
    isLoading,
  };
}
